package com.boge.ai.tool;

import com.boge.ai.bean.MilvusEntity;
import com.boge.ai.bean.UserInfo;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.tool.execution.ToolCallResultConverter;
import org.springframework.ai.tool.execution.ToolExecutionExceptionProcessor;

import java.lang.reflect.Type;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class Tools {

    @Tool(description = "计算两数之和")
    int add(@ToolParam(description = "第一个加数") int a, @ToolParam(description = "第二个加数") int b) {
        System.out.println("add方法执行了...");
        return a + b;
    }

    int add1( int a,  int b) {
        System.out.println("add1方法执行了...");
        return a + b;
    }

    @Tool(description = "病人患病信息检索", returnDirect = true,resultConverter = CustomConverter.class)
    UserInfo getUser(@ToolParam(description = "用户ID") Long id, ToolContext context) throws SQLException {
        Object tenantId = context.getContext().get("tenantId");
        System.out.println(tenantId);
        if (id ==1){
            throw new SQLException("数据库异常信息");
        }
        return new UserInfo(id,"李四");
    }


    @Tool(description = "获取用户信息")
    String vectorSearch(@ToolParam(description = "患者的问题") String msg, ToolContext context){
        // 1.获取查询的向量信息
        List<List<Float>> floats = new ArrayList<>();
        EmbeddingModel embeddingModel = (EmbeddingModel) context.getContext().get("embeddingModel");
        MilvusServiceClient client = (MilvusServiceClient) context.getContext().get("client");
        float[] embeddings = embeddingModel.embed("白癜风怎么治疗");
        // float[] embeddings 转换为 List<Float>
        List<Float> embeddingList = new ArrayList<>(embeddings.length);
        for (float f : embeddings) {
            embeddingList.add(f);
        }
        floats.add(embeddingList);
        // 封装向量数据库检索的信息
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(MilvusEntity.COLLECTION_NAME)
                .withMetricType(MetricType.L2)// 使用 L2 距离作为相似度度量
                .withTopK(3) // 返回最接近的前3个结果
                .withVectors(floats)
                .withVectorFieldName(MilvusEntity.Field.FEATURE) // 向量字段名
                .withOutFields(Arrays.asList(MilvusEntity.Field.ID,MilvusEntity.Field.OUTPUT,MilvusEntity.Field.TEXT)) // 需要返回的字段
                .build();
        SearchResults data = client.search(searchParam).getData();
        // 3.处理响应的信息
        SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());
        List<QueryResultsWrapper.RowRecord> rowRecords = resultsWrapper.getRowRecords();
        StringBuilder stringBuilder = new StringBuilder();
        if(rowRecords != null && !rowRecords.isEmpty()){
            stringBuilder.append(rowRecords);
        }
        return stringBuilder.toString();
    }

}

